# -*- coding: utf-8 -*-
# !/usr/bin/env python
"""
-------------------------------------------------
   File Name：     test
   Description :   
   Author :       lth
   date：          2023/1/3
-------------------------------------------------
   Change Activity:
                   2023/1/3 13:20: create this script
-------------------------------------------------
"""
__author__ = 'lth'

import numpy as np
import torch
import torchvision.utils
from PIL import Image

from config import GetConfig
from model import Generator
from utils import mixing_noise, denormalize, lerp
torch.cuda.is_available()

class Inference:
    def __init__(self):
        self.args = GetConfig()

        print(f"-----------{self.args.project_name}-------------")
        use_cuda = self.args.use_cuda and torch.cuda.is_available()
        self.device = torch.device("cuda" if use_cuda else "cpu")
        self.generator_ema = Generator(size=256,
                                       style_dim=512,
                                       n_mlp=8,
                                       channel_multiplier=2,
                                       blur_kernel=[1, 3, 3, 1],
                                       lr_mlp=0.01, )

        self.generator_ema.eval()
        if use_cuda:
            self.generator_ema = torch.nn.DataParallel(self.generator_ema, device_ids=range(torch.cuda.device_count()))
        self.generator_ema.load_state_dict(torch.load("weights/generate_best.pth")["model_state_dict"])
        self.sampling = "full"

    @torch.no_grad()
    def test(self):
        self.generator_ema.eval()
        latent = mixing_noise(4, 512, 0.5, self.device)
        fake_img, _ = self.generator_ema(latent, return_latents=False)
        self.get_image(fake_img, fake_img, save_path="./results/get_from_generate_ema.jpg")
        torchvision.utils.save_image(fake_img, "results/gema.png", nrow=2, normalize=True, value_range=(-1, 1))

    @torch.no_grad()
    def lerp_show(self):

        self.generator_ema.eval()

        count = 8

        noise = self.generator_ema.module.make_noise()
        inputs = torch.randn([6, 512], device=self.device)
        if self.sampling == "full":
            lerp_t = torch.rand(count, device=self.device)
        else:
            lerp_t = torch.zeros(count, device=self.device)
        lerp_t = torch.sort(lerp_t)[0]
        lerp_t2 = torch.arange(0, 1, 0.125, device=self.device)
        print(lerp_t)
        latent = self.generator_ema.module.get_latent(inputs)
        latent_t0, latent_t1, latent_t2, latent_t3, latent_t4, latent_t5 = latent[0], latent[1], latent[2], latent[3], \
                                                                           latent[4], latent[5]
        latent_e0 = lerp(latent_t1, latent_t0, lerp_t[:, None])
        latent_e1 = lerp(latent_t2, latent_t0, lerp_t[:, None])
        latent_e2 = lerp(latent_t3, latent_t0, lerp_t[:, None])
        latent_e3 = lerp(latent_t4, latent_t0, lerp_t[:, None])
        latent_e4 = lerp(latent_t5, latent_t0, lerp_t[:, None])

        latent_e = torch.cat([latent_e0, latent_e1, latent_e2, latent_e3, latent_e4], 0)

        image, _ = self.generator_ema(styles=[latent_e], input_is_latent=True, noise=noise)
        torchvision.utils.save_image(image, "./results/lerp_show.jpg", nrow=count,
                                     normalize=True, value_range=(-1, 1))

    @staticmethod
    def get_image(image, photo, save_path="./results/get_from_generate.jpg"):
        assert len(image) >= 4, "image batchsize should greater than 4, or this function can not work "

        image = torch.clamp(image,-1,1)
        photo = torch.clamp(photo,-1,1)

        output = (denormalize(image.permute((0, 2, 3, 1)).detach().to('cpu').numpy()) * 255).astype('uint8')
        photo = (denormalize(photo.permute((0, 2, 3, 1)).detach().to('cpu').numpy()) * 255).astype('uint8')

        # output = np.clip(output + 0.5, 0, 255).astype(np.uint8)
        # photo = np.clip(photo + 0.5, 0, 255).astype(np.uint8)

        output0 = Image.fromarray(output[0]).convert('RGB')
        photo0 = Image.fromarray(photo[0]).convert('RGB')

        output1 = Image.fromarray(output[1]).convert('RGB')
        photo1 = Image.fromarray(photo[1]).convert('RGB')

        output2 = Image.fromarray(output[2]).convert('RGB')
        photo2 = Image.fromarray(photo[2]).convert('RGB')

        output3 = Image.fromarray(output[3]).convert('RGB')
        photo3 = Image.fromarray(photo[3]).convert('RGB')

        width = output3.width
        height = output3.height

        target = Image.new('RGB', (width + width, 4 * height), (255, 255, 255))

        target.paste(output0, (0, 0, width, height))
        target.paste(photo0, (width, 0, 2 * width, height))

        target.paste(output1, (0, height, width, 2 * height))
        target.paste(photo1, (width, height, 2 * width, 2 * height))

        target.paste(output2, (0, height * 2, width, 3 * height))
        target.paste(photo2, (width, height * 2, width * 2, height * 3))

        target.paste(output3, (0, height * 3, width, 4 * height))
        target.paste(photo3, (width, height * 3, 2 * width, height * 4))

        target.save(save_path)


if __name__ == "__main__":
    model = Inference()
    # model.lerp_show()
    model.test()
